import numpy as np
import os
from collections import Counter


class NaiveBayes:
    def __init__(self):
        self.class_probs = {}
        self.word_probs = {}

    def train(self, data):
        class_counts = Counter([d[0] for d in data])
        total_data = len(data)
        self.class_probs = {cls: count / total_data for cls, count in class_counts.items()}

        word_counts = {cls: Counter() for cls in class_counts}
        total_words_in_class = {cls: 0 for cls in class_counts}

        for label, text in data:
            for word in text.split():
                word_counts[label][word] += 1
                total_words_in_class[label] += 1

        self.word_probs = {cls: {word: (count + 1) / (total_words_in_class[cls] + len(word_counts[cls]))
                                 for word, count in word_count.items()}
                           for cls, word_count in word_counts.items()}

    def predict(self, text):
        class_scores = {}
        for cls, prob in self.class_probs.items():
            score = np.log(prob)
            for word in text.split():
                if word in self.word_probs[cls]:
                    score += np.log(self.word_probs[cls][word])
            class_scores[cls] = score
        return max(class_scores, key=class_scores.get)


# Example usage
train_data = [('sci.space', 'NASA launches a satellite'),
              ('rec.autos', 'New car models are out')]
model = NaiveBayes()
model.train(train_data)
prediction = model.predict('NASA satellite')
print(prediction)
